from stage1 import generate_heatmap
from stage2 import predict as grading
from PIL import Image
import cv2
import numpy as np
import os
from segment import segment


def get_roi(img, seg):
    """
    First step: obtain the roi
    :param img: numpy array
    :param seg: numpy array
    :return:
    """
    heatmap_on_image, roi, only_heatmap_on_seg, nerve_gradient_on_img = generate_heatmap(
        img, seg)
    # roi will be used in stage2,
    # the only_heatmap_on_seg, nerve_gradient_on_img are used to visualize
    return heatmap_on_image, roi, only_heatmap_on_seg, nerve_gradient_on_img


def tortuosity_grading(img, seg):
    """
    Choose one to show/visualize as you desired
    Second step: obtain the grading results
    :param img: numpy array
    :param seg: numpy array
    :return: 4 items
        "only_heatmap_on_seg": the heatmap on segmentation
        "nerve_gradient_on_img": the heatmap on nerve and the orginal image
        "pred_tortuosity_level": int value, the predicted tortuosity level
        "pred_levels_probabilities": float values, the predicted probabilities of each level (0, 1, 2, 3)
        "heatmap_on_image": show heatmap on original image
    """
    heatmap_on_image, roi, only_heatmap_on_seg, nerve_gradient_on_img = get_roi(
        img, seg)
    # 'roi': <class 'numpy.ndarray'>
    pred_tortuosity_level, pred_levels_probabilities = grading(img, seg, roi)
    pred_tortuosity_level = pred_tortuosity_level[0]
    pred_levels_probabilities = pred_levels_probabilities[0]
    return heatmap_on_image, only_heatmap_on_seg, nerve_gradient_on_img, pred_tortuosity_level, pred_levels_probabilities


if __name__ == '__main__':
    print(os.getcwd())
    # usage example
    img_file = "./test_image/original_img.jpg"
    # seg_file = "./test_image/segmentation.png"
    img = cv2.imread(img_file, flags=-1)
    # seg = cv2.imread(seg_file, flags=-1)

    # get the segmentation first based on the CornealNerveSegDemo
    seg = segment(img)

    # then perform torutosity grading
    heatmap_on_image, only_heatmap_on_seg, nerve_gradient_on_img, pred_tortuosity_level, pred_levels_probabilities = tortuosity_grading(
        img, seg)
    print(pred_tortuosity_level)
    print(pred_levels_probabilities)
